"""
LSTM预测交易策略
策略逻辑：
1. 使用LSTM模型预测未来价格
2. 结合预测结果生成交易信号
3. 动态调整仓位大小
"""

import torch
import torch.nn as nn
import numpy as np
import pandas as pd
from sklearn.preprocessing import MinMaxScaler

class LSTMPredictor(nn.Module):
    def __init__(self, input_size=1, hidden_size=50, output_size=1):
        super().__init__()
        self.hidden_size = hidden_size
        self.lstm = nn.LSTM(input_size, hidden_size, batch_first=True)
        self.linear = nn.Linear(hidden_size, output_size)
        
    def forward(self, x):
        out, _ = self.lstm(x)
        out = self.linear(out[:, -1, :])
        return out

class LSTMTradingStrategy:
    def __init__(self, ticker, lookback=60, forecast_horizon=5):
        self.ticker = ticker
        self.lookback = lookback
        self.forecast_horizon = forecast_horizon
        self.scaler = MinMaxScaler(feature_range=(-1, 1))
        self.model = self.init_model()
        
    def init_model(self):
        """初始化LSTM模型"""
        model = LSTMPredictor()
        # 这里应该加载预训练权重
        return model
    
    def preprocess_data(self, prices):
        """预处理价格数据"""
        scaled = self.scaler.fit_transform(prices.reshape(-1, 1))
        sequences = []
        for i in range(len(scaled) - self.lookback - self.forecast_horizon + 1):
            seq = scaled[i:i+self.lookback]
            label = scaled[i+self.lookback:i+self.lookback+self.forecast_horizon]
            sequences.append((seq, label))
        return sequences
    
    def predict(self, recent_prices):
        """预测未来价格"""
        self.model.eval()
        with torch.no_grad():
            seq = torch.FloatTensor(recent_prices[-self.lookback:]).unsqueeze(0)
            pred = self.model(seq).numpy()
        return self.scaler.inverse_transform(pred)[0]
    
    def generate_signal(self, current_price, recent_prices):
        """生成交易信号"""
        if len(recent_prices) < self.lookback + self.forecast_horizon:
            return 0
            
        # 预测未来5天价格
        preds = self.predict(recent_prices)
        avg_pred = np.mean(preds)
        
        # 计算预期收益率
        expected_return = (avg_pred - current_price) / current_price
        
        # 根据预期收益率确定仓位
        if expected_return > 0.05:  # 预期上涨5%以上
            return min(1.0, expected_return * 2)  # 杠杆不超过2倍
        elif expected_return < -0.05:  # 预期下跌5%以上
            return max(-1.0, expected_return * 2)
        else:
            return 0

if __name__ == '__main__':
    strategy = LSTMTradingStrategy('AAPL')
    
    # 模拟价格数据
    np.random.seed(42)
    test_prices = np.cumsum(np.random.randn(100)) + 100
    
    # 生成信号
    signal = strategy.generate_signal(test_prices[-1], test_prices)
    print(f"预测信号强度: {signal:.2f}")